#!/usr/bin/env python3
# coding: utf-8
# Copyright 2024 Huawei Technologies Co., Ltd
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
# ===========================================================================
import os
import json
import platform

from ansible.module_utils.basic import AnsibleModule
from ansible.module_utils import common_info, common_utils, venv_installer
from ansible.module_utils.common_info import DeployStatus, get_os_and_arch


class GrafanaInstaller:
    def __init__(self):
        self.module = AnsibleModule(
            argument_spec=dict(
                tags=dict(type="list", required=True),
                resources_dir=dict(type="str", required=True),
                ip_address=dict(type="str", required=True)
            )
        )
        self.run_tags = self.module.params["tags"]
        self.arch = platform.machine()
        self.resources_dir = os.path.expanduser(self.module.params["resources_dir"])
        self.ip_address = self.module.params["ip_address"]
        self.provision_path = "/etc/grafana/provisioning"
        self.messages = []

    def _module_failed(self):
        return self.module.fail_json(msg="\n".join(self.messages), rc=1, changed=False)

    def _module_success(self):
        return self.module.exit_json(msg="Install grafana success.", rc=0, changed=True)

    def _find_grafana_pkg(self, pattern):
        pkgs, msgs = common_utils.find_files(os.path.join(self.resources_dir, "Grafana*"), pattern)
        self.messages.extend(msgs)
        if not pkgs:
            if "auto" in self.run_tags:
                self.module.exit_json(std_out="can not find grafana package, grafana install skipped", rc=0,
                                      result={DeployStatus.DEPLOY_STATUS: DeployStatus.SKIP},
                                      changed=False)
            else:
                self.messages.append("can not find grafana package.")
                self._module_failed()

        return pkgs
    
    def _get_config(self):
        config_path = self._find_grafana_pkg("npu_monitor_config.json")[0]
        with open(config_path, 'r') as f:
            data = json.load(f)
            self.config = data.get("grafana", {})
            self.prometheus_port = data.get("prometheus", {}).get("port", 9090)
    
    def update_grafana_ini(self):
        config_map = {
            "http_port":self.config.get("port", 3000),
            "defaul_language":self.config.get("language", "zh-Hans")
        }
        with open("/etc/grafana/grafana.ini", "r") as f:
            lines = f.readlines()
        for key, vaule in config_map.items():
            pattern = ";{} =".format(key)
            for i in range(len(lines)):
                if pattern in lines[i]:
                    lines[i] = "{} = {}".format(key, vaule)
        with open("/etc/grafana/grafana.ini", "w") as f:
            f.writelines(lines)

    def install_pkg(self, pkg_type):
        pkgs = self._find_grafana_pkg("*{}.{}".format(self.arch, pkg_type))
        for pkg in pkgs:
            if pkg_type == "rpm":
                 self.module.run_command(["rpm", "-ivh", "--nodeps", pkg])
            else:
                self.module.run_command(["dpkg", "-i", pkg])
        self.messages.append("install pkg success")
    
    def create_datasource_yaml(self):
        yaml_path = os.path.join(self.provision_path, "datasources/prometheus.yaml")
        content = [
            "apiVersion: 1",
            "datasources:",
            "  - name: prometheus",
            "    type: prometheus",
            "    access: proxy",
            "    orgId: 1",
            "    uid: my_unique_uid",
            "    url: http://{}:{}".format(self.ip_address, self.prometheus_port)
        ]
        with open(yaml_path, 'w') as f:
            f.write("\n".join(content))
    
    def create_dashboards_config(self):
        dashboards_config = self._find_grafana_pkg("grafana_dashboards_config*.json")[0]
        dashboards_config_dir = os.path.join(self.provision_path, "grafana_dashboards_config")
        if os.path.exists(dashboards_config_dir):
            if not os.path.isdir(dashboards_config_dir):
                self.messages.append("{} is not dir".format(dashboards_config_dir))
                self._module_failed()
        else:
            self.module.run_command(["mkdir", "-p", dashboards_config_dir])
        with open(dashboards_config, 'r') as f:
            data = json.load(f)
        data["templating"]["list"].append({
            "hide": 0,
            "label": "datasource",
            "name": "DS_PROMETHEUS",
            "options": [],
            "query": "prometheus",
            "refresh": 1,
            "regex": "",
            "type": "datasource"
        })
        with open(os.path.join(dashboards_config_dir, "grafana_dashboards_config.json"), 'w') as f:
            data = json.dump(data, f, indent=4)
        return dashboards_config_dir
    
    def create_dashboards_yaml(self):
        dashboards_config_dir = self.create_dashboards_config()
        yaml_path = os.path.join(self.provision_path, "dashboards/npu_monitor.yaml")
        content = [
            "apiVersion: 1",
            "providers:",
            "  - name: prometheus",
            "    type: file",
            "    updateIntervalSeconds: 30",
            "    options:",
            "      path: {}".format(dashboards_config_dir),
            "      foldersFromFilesStructure: true"
        ]
        with open(yaml_path, 'w') as f:
            f.write("\n".join(content))
    
    def setup_service(self):
        self.module.run_command(["systemctl", "enable", "grafana-server"])
        self.module.run_command(["systemctl", "start", "grafana-server"])

    def run(self):
        try:
            self._get_config()
            os_and_arch = get_os_and_arch()
            if os_and_arch.startswith(("Ubuntu", "Debian", "veLinux")):
                self.install_pkg("deb")
            else:
                self.install_pkg("rpm")
            self.update_grafana_ini()
            self.create_datasource_yaml()
            self.create_dashboards_yaml()
            self.setup_service()
        except Exception as e:
            self.messages.append(str(e))
            self._module_failed()
        self._module_success()


def main():
    GrafanaInstaller().run()


if __name__ == "__main__":
    main()
